package top.continew.ai.service.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.io.FileUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.PostMapping;
import top.continew.ai.service.ParseDocmentService;
import top.continew.ai.service.VectorStoreService;
import top.continew.ai.util.parse.BaseParse;
import top.continew.ai.util.parse.ParseFactory;

import java.io.File;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 解析文档 并存入到向量数据库中
 *
 * @Author： Lonni
 * @Date： 2025/8/21  15:13
 */
@Service
@Slf4j
public class ParseDocmentServiceImpl implements ParseDocmentService {

    private final VectorStoreService vectorStoreService;

    // 从配置文件注入分片大小和滑动窗口大小
    @Value("${rag.chunk.size:1000}")
    private int chunkSize;

    @Value("${rag.slide.window.size:200}")
    private int slideWindowSize;

    public ParseDocmentServiceImpl(VectorStoreService vectorStoreService) {
        this.vectorStoreService = vectorStoreService;
    }


    /**
     * TODO 测试加载文档到知识库
     */

    @PostMapping
    public void loadFile() {
        String path = "classpath:/docment";
        List<File> files = FileUtil.loopFiles(path);
        if (CollectionUtil.isEmpty(files)) {
            log.info("没有要加载的文件");
            return;
        }

        files.forEach(p->{
            parseDocment(p.getAbsolutePath(),p.getName());

        });

        log.info("解析文档成功");


    }


    /**
     * 解析文档
     *
     * @param filePath 文件路径
     * @param fileName 文件名
     */
    @Override
    public void parseDocment(String filePath, String fileName) {
        try {
            BaseParse parse = ParseFactory.me.getParse(fileName);
            List<Document> documents = parse.renderDocment(filePath, fileName);

            if (CollectionUtil.isEmpty(documents)) {
                log.info("上传的文档为空,将删除且不做解析处理");
                FileUtil.del(filePath);
                return;
            }

            // 使用自定义的分片方法处理文档
            List<Document> splitDocuments = splitDocumentsWithSlidingWindow(documents, fileName);

            // 存入向量数据库
            vectorStoreService.addDocuments(splitDocuments);

        } catch (Exception e) {
            log.error("解析文档时发生错误", e);
        }
    }

    /**
     * 使用 TokenTextSplitter 固定大小分片 + 滑动窗口算法分片
     * 当长度大于指定大小时，继续使用滑动窗口算法分片
     *
     * @param documents 原始文档列表
     * @return 分片后的文档列表
     */
    private List<Document> splitDocumentsWithSlidingWindow(List<Document> documents, String fileName) {
        List<Document> resultDocuments = new ArrayList<>();

        // 首先使用 TokenTextSplitter 进行基础分片
        ChineseTokenTextSplitter tokenSplitter = new ChineseTokenTextSplitter(chunkSize, slideWindowSize, 5, 20000, true);
        List<Document> tokenSplitDocuments = tokenSplitter.apply(documents);

        // 对每个分片检查是否需要进一步滑动窗口分片
        for (Document document : tokenSplitDocuments) {
            document.getMetadata().put("fileName", fileName);
            String content = document.getText();
            // 如果内容长度小于等于分片大小，直接添加
            if (!StringUtils.hasText(content) || content.length() <= chunkSize) {
                resultDocuments.add(document);
                continue;
            }

            // 如果内容长度大于分片大小，使用滑动窗口进一步分片
            List<String> slidingChunks = splitWithSlidingWindow(content);

            // 为每个滑动分片创建新的 Document 对象
            for (int i = 0; i < slidingChunks.size(); i++) {
                String chunk = slidingChunks.get(i);
                Document chunkDocument = new Document(chunk, document.getMetadata());
                // 可以添加额外的元数据来标识分片信息
                Map<String, Object> metadata = chunkDocument.getMetadata();
                metadata.put("original_chunk_index", tokenSplitDocuments.indexOf(document));
                metadata.put("sliding_chunk_index", i);
                resultDocuments.add(chunkDocument);
            }
        }

        return resultDocuments;
    }

    /**
     * 使用滑动窗口算法对长文本进行分片
     *
     * @param content 需要分片的文本内容
     * @return 分片后的文本列表
     */
    private List<String> splitWithSlidingWindow(String content) {
        List<String> chunks = new ArrayList<>();

        // 如果内容长度小于分片大小，直接返回整个内容
        if (content.length() <= chunkSize) {
            chunks.add(content);
            return chunks;
        }

        // 应用滑动窗口分片
        int start = 0;
        while (start < content.length()) {
            int end = Math.min(start + chunkSize, content.length());
            chunks.add(content.substring(start, end));

            // 移动滑动窗口，最后一片不需要再滑动
            if (end == content.length()) {
                break;
            }

            // 滑动距离 = 分片大小 - 滑动窗口大小（重叠部分）
            start = Math.min(start + (chunkSize - slideWindowSize), content.length() - chunkSize);
        }

        return chunks;
    }
}
